import struct
from pathlib import Path

import pandas as pd

from ..base import BaseReader
from ...exceptions import TdxFileNotFoundException


class TdxExHqDailyBarReader(BaseReader):
    """
    读取通达信扩展日线数据
    """

    def parse_data_by_file(self, fname):
        """

        :param fname:
        :return:
        """
        if not Path(fname).is_file():
            raise TdxFileNotFoundException('no tdx kline data, pleaes check path %s', fname)

        with open(fname, 'rb') as f:
            return self.unpack_records('<IffffIIf', f.read())

    def get_df(self, code_or_file, exchange=None):
        """

        :param code_or_file:
        :param exchange:
        :return:
        """
        # 只传入了一个参数
        data = [self._df_convert(row) for row in self.parse_data_by_file(code_or_file)]
        columns = ('date', 'open', 'high', 'low', 'close', 'amount', 'volume', 'jiesuan', 'hk_stock_amount')

        df = pd.DataFrame(data=data, columns=columns)
        df.index = pd.to_datetime(df.date)

        return df[['open', 'high', 'low', 'close', 'amount', 'volume', 'jiesuan', 'hk_stock_amount']]

    @staticmethod
    def _df_convert(row):
        """

        :param row:
        :return:
        """
        t_date = str(row[0])
        date_str = t_date[:4] + '-' + t_date[4:6] + '-' + t_date[6:]

        (hk_stock_amount,) = struct.unpack('<f', struct.pack('<I', row[5]))
        new_row = (date_str, row[1], row[2], row[3], row[4], row[5], row[6], row[7], hk_stock_amount)

        return new_row
